# coding=utf-8
# 学习单位  : 郑州大学
# @Author  : 铭同学
# @Time    : 2021/10/13 8:47
# @Software: PyCharm

# 多分类 Softmax
import torch
from torchvision import transforms       # 图像处理
from torchvision import datasets         # 数据集
from torch.utils.data import DataLoader  # 数据载入
import torch.nn.functional as F          # 函数包--激活函数使用（relu()）
import torch.optim as optim              # 优化器

# 1. Prepare Dataset
batch_size = 64
transform = transforms.Compose([
    transforms.ToTensor(),  # 将PIL Image转换为Tensor（这时通道在前，通道(c)*宽度(w)*高度(h)）
    transforms.Normalize((0.1307, ), (0.3081, ))    # 对数据标准化转换为0-1分布
])

# 准备数据集
train_dataset = datasets.MNIST(root='../dataset/mnist/' #数据路径
                               ,train=True  # 作为训练集
                               ,download=True   # 如果路径无数据进行下载
                               ,transform=transform
                               )
# 批量载入数据集
train_loader = DataLoader(train_dataset #数据集
                          ,shuffle=True # 打乱数据
                          ,batch_size=batch_size    #批量大小
                          )

# 准备数据集
test_dataset = datasets.MNIST(root='../dataset/mnist/' #数据路径
                               ,train=False
                               ,download=True   # 如果路径无数据进行下载
                               ,transform=transform
                               )
# 批量载入数据集
test_loader = DataLoader(test_dataset #数据集
                          ,shuffle=False # 不打乱数据，有利于结果观察
                          ,batch_size=batch_size    #批量大小
                          )
# 2. Design Model
# 模型类--全连接神经网络（全部都是线性层串联）
class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.l1 = torch.nn.Linear(784, 512)
        self.l2 = torch.nn.Linear(512, 256)
        self.l3 = torch.nn.Linear(256, 128)
        self.l4 = torch.nn.Linear(128, 64)
        self.l5 = torch.nn.Linear(64, 10)

    def forward(self, x):   # 必须实现同名forward，注意拼写错误
        x = x.view(-1, 784) # 转换tensor数据维度（tensor.view()功能相当于numpy.reshape()）
        # 这里的激活函数（非线性函数）使用relu
        x = F.relu(self.l1(x))
        x = F.relu(self.l2(x))
        x = F.relu(self.l3(x))
        x = F.relu(self.l4(x))
        return self.l5(x)   # 最后一层不需要激活函数，直接输出线性结果（交叉熵接收的就是线性结果）
# 实例化模型
model = Net()

# 3. Construct Loss and Optimizer
criterion = torch.nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)    # 加一个动（冲）量（其作用就是摆脱局部最优，奔向全局最优）

# 4. Train and Test
def train(epoch):
    running_loss = 0.0
    # TypeError: __array__() takes 1 positional argument but 2 were given
    # 解决办法，降低pillow版本到8.2即可，pip install pillow==8.2
    for batch_idx, data in enumerate(train_loader, 0): # 这一行读数据失败--已经解决
        inputs, target = data
        optimizer.zero_grad()   # 梯度归零，避免累加

        # forward + backward + update
        outputs = model(inputs)
        loss = criterion(outputs, target)
        loss.backward()
        optimizer.step()

        running_loss += loss.item() # item()用于取出值本身，对每个批次的损失进行求和
        # 每300个批次输出一次
        if batch_idx % 300 == 299:
            print('[%d, %5d] loss: %.3f' % (epoch + 1, batch_idx + 1, running_loss / 300))
            running_loss = 0.0 # 损失清零

def test():
    correct = 0
    total = 0
    # 测试时不计算梯度（求梯度是很贵的事情，非必要即可不求）
    with torch.no_grad():
        for data in test_loader:
            images, labels = data
            outputs = model(images)
            _, predicted = torch.max(outputs.data, dim=1)   # 按照每一行求最大值（就是最大概率），输出两个结果（数值本身和对应下标）我们只要下标predicted(也就是对应的分类)
            total += labels.size(0) # 计算总数
            # 计算预测正确的个数
            correct += (predicted == labels).sum().item()

    print('Accuracy on test set: %.2f%% [%d / %d]' % ((100 * correct / total), correct, total))


if __name__ == '__main__':
    pass
    # print(type(train_loader))
    # print(train_loader)
    # for data in train_loader:
    #     print(data)
    # 训练和测试进行分装后，主函数更简洁
    for epoch in range(10):
        train(epoch)
        test()


'''控制台结果，训练到97就是一个瓶颈了

[1,   300] loss: 2.133
[1,   600] loss: 0.745
[1,   900] loss: 0.411
Accuracy on test set: 88 %
[2,   300] loss: 0.312
[2,   600] loss: 0.259
[2,   900] loss: 0.223
Accuracy on test set: 94 %
[3,   300] loss: 0.188
[3,   600] loss: 0.164
[3,   900] loss: 0.146
Accuracy on test set: 95 %
[4,   300] loss: 0.131
[4,   600] loss: 0.122
[4,   900] loss: 0.109
Accuracy on test set: 96 %
[5,   300] loss: 0.096
[5,   600] loss: 0.097
[5,   900] loss: 0.089
Accuracy on test set: 96 %
[6,   300] loss: 0.076
[6,   600] loss: 0.075
[6,   900] loss: 0.075
Accuracy on test set: 97 %
[7,   300] loss: 0.056
[7,   600] loss: 0.063
[7,   900] loss: 0.065
Accuracy on test set: 97 %
[8,   300] loss: 0.044
[8,   600] loss: 0.054
[8,   900] loss: 0.053
Accuracy on test set: 97 %
[9,   300] loss: 0.037
[9,   600] loss: 0.042
[9,   900] loss: 0.042
Accuracy on test set: 97 %
[10,   300] loss: 0.034
[10,   600] loss: 0.032
[10,   900] loss: 0.033
Accuracy on test set: 97 %


'''